884C - Bertown Subway - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100005;
int a[MAXN], b[MAXN];
bool vis[MAXN];
int dfs(int s,int next,int res)
{
	if(next == s)
    {
        return res;
    }
	if(!vis[next])
	{
		vis[next] = true;
		return dfs(s,a[next],res+1);
	}
	return res;
}
int main()
{
	int n;
	cin>>n;
	for(int i = 1; i <= n; i++)
    {
        cin>>a[i];
    } 
	memset(vis,false,sizeof(vis));
	int cnt = 0;
	for (int i = 1; i <= n; i++)
    {
        if(!vis[i])
		{
			vis[i] = true;
			b[cnt++] = dfs(i,a[i],1);
		}
    }
	sort(b,b+cnt);
	long long ans = 0;
	for(int i = 0; i < cnt-2; i++)
    {
        ans += b[i]*b[i];
    }
	ans += (long long)(b[cnt-2]+b[cnt-1])*(b[cnt-2]+b[cnt-1]);
	cout<<ans<<endl;
	return 0;
}

/**
* @runId: 321747
* @language: GNU G++17
* @author: 于沐阳2023
* @submitTime: 2024-02-18 10:53:53
*/